# -*- coding: utf-8 -*-
"""
Created on Tue Nov 24 17:02:27 2020

@author: Administrator
"""

import numpy as np
import math
import seaborn as sns
import pandas as pd
import copy



def Gauss(X, u, std):
    """
    gaussian funtion
    """
    # One-dimensional Gaussian distribution 
    result = (1 / (math.sqrt(2 * math.pi) * std)) * \
             np.exp(-1 * (X - u) * (X - u) / (2 * std**2))
    return result

def E_step(data, num_kernal, parameters):
    """
    expection,
    input:
        data, dataset
        num_kernal, number of kernal
        parameters, dictionary use to record probablity, mean, std
    """
    # gamma jk. the probability of jth observation comes from kth gaussian function
    gamma = {}
    for i in range(num_kernal):
        # ak * gaussian(yi | theta k)
        gamma[i] = parameters[i]['a'] * Gauss(data, parameters[i]['u'], parameters[i]['std'])
    
    _sum = sum(gamma.values())
    # the response of model k to observation data y
    for i in range(num_kernal):
        gamma[i] = gamma[i] / _sum
    
    return gamma

def M_step(data, num_kernal, old_para, gamma):
    """
    M step, updated the parameters
    input:
        data, dataset
        num_kernal, number of kernal
        old_para, the parameter in t-times.
        gamma, the response of modek k to observation y
    """
    parameters = copy.deepcopy(old_para)
    for i in range(num_kernal):
        #update mean, std, ak
        parameters[i]['u'] = np.dot(gamma[i], data) / np.sum(gamma[i])
        parameters[i]['std'] = np.sqrt(np.dot(gamma[i], (data-parameters[i]['u'])**2) / np.sum(gamma[i]))
        parameters[i]['a'] = np.sum(gamma[i]) / len(gamma[i])
        
    return parameters

def set_up_parameters(data, kernal):
    """
    set initial parameters
    """
    from collections import defaultdict
    
    parameters = defaultdict(dict)

    data = np.sort(data)
    _len = len(data)
    #split data into k groupby 
    data = np.array_split(data, kernal)
    
    #inistal parameters based upon the data
    for i in range(kernal):
        parameters[i]['u'] = data[i].mean()
        parameters[i]['std'] = data[i].std()
        parameters[i]['a'] = len(data[i]) / _len
        
    return parameters
    

def check_error(parameters, kernal, data):
    """
    Simulate data with newest parameter
    then calculated the error
    input:
        parameters: the parameters from EM algorithm
        kernal: number of kernel
        data: raw data
    """
    
    first = True
    for i in range(kernal):
        if first:
            first = False
            res = np.random.normal(parameters[i]['u'], parameters[i]['std'], int(parameters[i]['a'] * len(data)))
        else:
            tmp = np.random.normal(parameters[i]['u'], parameters[i]['std'], int(parameters[i]['a'] * len(data)))
            res = np.concatenate((res, tmp))
    # sort raw data
    sort_raw = np.sort(data)
    # sort data which simulate with EM algorithm's parameter
    sort_res = np.sort(res)
    # make sure their length is same 
    sort_raw = sort_raw[:len(sort_res)]
    # calcluated error
    error = abs(sort_raw - sort_res).sum()
    
    return error


if __name__ == '__main__':
    # dataset gaussian distreibution
    data0 = np.random.normal(10, 1, 300)
    data1 = np.random.normal(20, 2, 700)
    data2 = np.random.normal(30, 1, 400)
    # the ak = (.21, .28, .5)
    data = np.concatenate((data0, data1, data2))
    
    # check If the kernel increases, will the error decrease
    # kernal from 2 to 6
    for i in range(2, 6):
        # set up parameters
        parameters = set_up_parameters(data, i)
        step = 0
        while step < 500:
            step += 1
            # E step get gamma jk, the response of model k to observation data y 
            gamma = E_step(data, i, parameters)
            # updated parameters 
            new_parameters = M_step(data, i, parameters, gamma)
            # if parameters' change less than 1e-4. stop the iteration
            old = pd.DataFrame(parameters)
            new = pd.DataFrame(new_parameters)
            parameters = new_parameters
            if max(abs(old - new)) < 1e-4:
                break
        
        print (i, check_error(parameters, i, data))